/**********************************************************************/
/*
      Title: randinference.ado
			Author: Robbie Dulin
			Created: Dec 6 2019
    	Description: Program that implements randomization inference
        using resampling from an input file.
*/
/**********************************************************************/


cap program drop randinference
program define randinference, rclass
  version 14.2
  syntax varlist(numeric max = 1) using/, reg(string) reps(integer) mergevar(varlist) mergetype(string) [hetero(string asis) interaction(string asis)]

  * make sure hetero and interaction options specified correctly
  local num_hetero : list sizeof local(hetero)
  local num_interact : list sizeof local(interaction)
  cap assert `num_interact' == `num_hetero'
  if _rc != 0 {
    di as red "The same number of variables must specified in the hetero() and interaction() options."
    exit _rc
  }

  * initialize lists of hetero and interaction vars
  if `num_hetero' >= 1 {
    tokenize `hetero'
    forval i = 1 / `num_hetero' {
      local hetero`i' = "``i''"
    }

    tokenize `interaction'
    forval i = 1 / `num_interact' {
      local interaction`i' = "``i''"
      // di "`interaction`i''"
    }
  }

  * checks using file is okay
  preserve
  quietly use "`using'", clear
  // check that file has enough permutations of the resampvar
  forval i = 1 / `reps' {
    cap confirm variable `varlist'`i'
    if _rc != 0 {
      di as red "There are not enough permutations of `varlist'. There must be at least `reps' permutations."
      exit _rc
    }
  }
  restore


  * run original regression
  di "Original regression:"
  `reg'
  local original_obs = e(N)

  * initialize temporary scalar names
  tempname beta_original success success_fraction
  scalar `beta_original' = _b[`varlist']

  // set counters
  scalar `success' = 0

  * temporary scalars for interaction term
  if `num_hetero' >= 1 {
    forval i = 1 / `num_hetero' {
      tempname beta_original_interact`i' success_interact`i' success_fraction_interact`i'
      scalar `beta_original_interact`i'' = _b[`interaction`i'']
      scalar `success_interact`i'' = 0
    }
  }

  * run regressions on resampled treatment var
  preserve
  // drop original treatment var
  drop `varlist'

  // merge in resampvar
  quietly merge `mergetype' `mergevar' using "`using'"
  cap assert _m == 3 if e(sample) == 1
  if _rc != 0 {
    di as red "Some in-sample observations do not have matches in the resampling file."
    exit
  }
  drop _m

  // run `reps' regressions
  di "Beginning `reps' regressions..."
  forvalues num = 1 / `reps' {
    rename `varlist'`num' `varlist'

    * create interactions if hetero and interaction options specified
    if `num_hetero' >= 1 {
      forval i = 1 / `num_interact' {
        drop `interaction`i''
        qui gen `interaction`i'' = `varlist' * `hetero`i''
      }
    }

    // if "`hetero'" != "" {
    //   forval i = 1 / `num_interact' {
    //     rename `interaction`i''`num' `interaction`i''
    //   }
    // }

    // regression on resampvar
    quietly `reg'
    cap assert e(N) == `original_obs'
    if _rc != 0 {
      di as red "Iterated regressions do not match original"
      exit
    }

    // calculate if magnitude of Beta-i greater than magnitude of Beta-original
    scalar `success' = `success' + `=abs(_b[`varlist']) >= abs(`beta_original')'

    rename `varlist' `varlist'`num'

    // do the same for interaction term
    if `num_hetero' >= 1 {
      forval i = 1 / `num_interact' {
        scalar `success_interact`i'' = `success_interact`i'' + `=abs(_b[`interaction`i'']) >= abs(`beta_original_interact`i'')'
        // rename `interaction`i'' `interaction`i''`num'
      }
    }

    if inlist(`num', 10, 25, 50) == 1 {
      di "Completed regression `num'"
    }

    if mod(`num', 100) == 0 {
      di "Completed regression `num'"
    }

  }

  restore
  scalar `success_fraction' = `success' / `reps'
  return scalar RIpval =  `success_fraction'
  local pval = `success_fraction'

  if `num_hetero' >= 1 {
    di "RI p-val `varlist' = `pval'"
    forval i = 1 / `num_hetero' {
      scalar `success_fraction_interact`i'' = `success_interact`i'' / `reps'
      return scalar RIpval_interact`i' =  `success_fraction_interact`i''
      local pval_interact`i' = `success_fraction_interact`i''
      di "RI p-val `interaction`i'' (interaction #`i') = `pval_interact`i''"
    }
  }

  else {
    di "RI p-val = `pval'"
  }

  end
